// Author: Ivan Chau
package uk.ac.cam.kcic2.rayTracer;

import java.awt.Color;
import java.awt.image.BufferedImage;
import java.io.File;
import java.io.IOException;

import javax.imageio.ImageIO;

import uk.ac.cam.kcic2.math.Params;
import uk.ac.cam.kcic2.math.Vector3D;
import uk.ac.cam.kcic2.scene.Colour;
import uk.ac.cam.kcic2.scene.ColourCoeffs;
import uk.ac.cam.kcic2.scene.IntersectionInfo;
import uk.ac.cam.kcic2.scene.LightSource;
import uk.ac.cam.kcic2.scene.Ray;
import uk.ac.cam.kcic2.scene.Scene;

public class RayTracer {

	protected final Scene scene;
	// assume the camera at origin in world coordinates
	protected final Vector3D camera = new Vector3D(0,0,0);
	// assume the screen centred at (0,0,d) in world coordinates
	protected final double screenFromCamera;
	protected final double screenHeight;
	protected final double screenWidth;
	// divide the screen into a grid of hPixelsNum*vPixelsNum
	protected final int hPixelsNum;
	protected final int vPixelsNum;
	// the offsets of the centre in each pixel from the bottom left corner
	protected final double pixelCentreX;
	protected final double pixelCentreY;
	// the max number of bounces
	protected final int maxBounces;
	protected final double focalLength = 1;

	public RayTracer(Scene sceneSettings, double screenDistance, 
			double width, double height, int hPixels, int vPixels, int bounceNum){
		scene = sceneSettings;
		screenFromCamera = screenDistance;
		screenHeight = height;
		screenWidth = width;
		hPixelsNum = hPixels;
		vPixelsNum = vPixels;
		pixelCentreX = screenWidth/hPixelsNum/2d;
		pixelCentreY = screenHeight/vPixelsNum/2d;
		maxBounces = bounceNum;
	}

	public boolean renderImage(String filename){
		BufferedImage sceneImage = new BufferedImage(hPixelsNum,vPixelsNum,
				BufferedImage.TYPE_INT_RGB);
		int[] hitMap = computeColourMap();

		// produce a RGB image from the hitMap
		sceneImage.setRGB(0, 0, hPixelsNum, vPixelsNum, hitMap, 0, hPixelsNum);
		// save the image on disk in the defined location
		try {
			File outputfile = new File(filename + ".png");
			ImageIO.write(sceneImage, "png", outputfile);
			return true;
		} catch (IOException e) {
			return false;
		}
	}

	protected int[] computeColourMap(){
		int[] hitMap = new int[hPixelsNum*vPixelsNum];

		// trace a ray through each pixel on the screen
		for (int row = 0; row < vPixelsNum; row++){
			for (int column = 0; column < hPixelsNum; column++){
				Vector3D direction = new Vector3D(new double[]{
						-screenWidth/2d + (2*column + 1)*pixelCentreX,
						screenHeight/2d - (2*row + 1)*pixelCentreY,
						screenFromCamera
				});
				Ray cameraRay = new Ray(camera,direction);
				// compute the hitMap
				hitMap[row*hPixelsNum + column] = trace(cameraRay, 0).toRGB();
			}
		}
		return hitMap;
	}

	private Colour trace(Ray ray, int bounceCount){		
		// find the closest intersection
		IntersectionInfo intersection = scene.getIntersection(ray);

		// if there is an intersection
		if (intersection != null){
			// compute the shading
			Colour colour = shade(intersection, ray, bounceCount);
			return colour;
		}
		else{
			if (bounceCount == 0){
				return new Colour(Color.WHITE.getRGB());
			}
			else{
				return new Colour(Color.BLACK.getRGB());
			}
		}
	}

	private Colour shade(IntersectionInfo intersection, Ray viewer, int bounceCount){
		ColourCoeffs objCoeffs = intersection.object.colourCoeffs;
		Vector3D normal = intersection.normal;
		Vector3D viewerFromPoint = viewer.origin.sub(intersection.position);
		viewerFromPoint = viewerFromPoint.div(Math.sqrt(viewerFromPoint.norm()));

		// compute the ambient shading
		double red = scene.ambient.x*objCoeffs.getAmbientCoeffs(Params.X);
		double green = scene.ambient.y*objCoeffs.getAmbientCoeffs(Params.Y);
		double blue = scene.ambient.z*objCoeffs.getAmbientCoeffs(Params.Z);

		// find the contribution from each light source
		for (LightSource light : scene.lights){
			if (!inShadow(intersection, light)){
				Colour lightColour = light.colour;
				Vector3D lightFromPoint = light.position.sub(intersection.position);
				lightFromPoint = lightFromPoint.div(Math.sqrt(lightFromPoint.norm()));

				double[] diffuseColour = diffuseShading(lightColour, objCoeffs, normal, lightFromPoint);
				double[] specularColour = specularShading(lightColour, objCoeffs, 
						normal, lightFromPoint, viewerFromPoint);

				red += diffuseColour[Params.X] + specularColour[Params.X];
				green += diffuseColour[Params.Y] + specularColour[Params.Y];
				blue += diffuseColour[Params.Z] + specularColour[Params.Z];
			}
		}

		// find any reflections
		if (objCoeffs.specularCoeffs.norm() > 0){	// specular surface
			if (bounceCount < maxBounces){
				Vector3D reflectionDirection = normal.mult(2*viewerFromPoint.dot(normal)).sub(viewerFromPoint);
				Ray reflectionRay = new Ray(intersection.position, reflectionDirection);
				Colour reflectedColour = trace(reflectionRay, bounceCount+1);

				double[] specularColour = specularShading(reflectedColour, objCoeffs, 
						normal, reflectionDirection, viewerFromPoint);
				red += specularColour[Params.X];
				green += specularColour[Params.Y];
				blue += specularColour[Params.Z];
				//red += reflectedColour.get(Params.X)*objCoeffs.getSpecularCoeffs(Params.X);
				//green += reflectedColour.get(Params.Y)*objCoeffs.getSpecularCoeffs(Params.Y);
				//blue += reflectedColour.get(Params.Z)*objCoeffs.getSpecularCoeffs(Params.Z);
			}
		}

		// find any refractions
		if (objCoeffs.transmittance.norm() > 0){	// transparent surface
			Vector3D refractedDirection = refract(viewerFromPoint.mult(-1), normal, 1/objCoeffs.refractiveIndex);
			Ray refractionRay = new Ray(intersection.position, refractedDirection);
			IntersectionInfo nextIntersection = scene.getIntersection(refractionRay);

			int internalBounces = 1;	
			while(nextIntersection != null && intersection.object == nextIntersection.object && 
					internalBounces + bounceCount < maxBounces){
				// ray travels inside the object				
				refractedDirection = refract(refractedDirection, nextIntersection.normal, 
						objCoeffs.refractiveIndex);
				refractionRay = new Ray(nextIntersection.position, refractedDirection);
				nextIntersection = scene.getIntersection(refractionRay);
				internalBounces++;	

			}

			if (nextIntersection != null && nextIntersection.object != intersection.object){
				// ray comes back out of the object
				Colour refractedColour = trace(refractionRay, internalBounces + bounceCount);

				red += refractedColour.get(Params.X)*Math.pow(objCoeffs.getTransmittance(Params.X), internalBounces);
				green += refractedColour.get(Params.Y)*Math.pow(objCoeffs.getTransmittance(Params.Y), internalBounces);
				blue += refractedColour.get(Params.Z)*Math.pow(objCoeffs.getTransmittance(Params.Z), internalBounces);
			}
		}




		return new Colour(new double[]{red, green, blue});
	}

	public static Vector3D refract(Vector3D incidentDirection, Vector3D normal, double refractiveIndicesRatio){

		double cosTheta1 = incidentDirection.dot(normal);
		double cos2Theta2 = 1 - refractiveIndicesRatio*refractiveIndicesRatio*(1 - cosTheta1*cosTheta1);

		if (cos2Theta2 < 0){ // Total internal reflection
			Vector3D viewerFromPoint = incidentDirection.mult(-1);
			Vector3D reflectionDirection = normal.mult(2*viewerFromPoint.dot(normal)).sub(viewerFromPoint);
			return reflectionDirection;
		}
		else{
			double value = refractiveIndicesRatio*cosTheta1 + Math.sqrt(cos2Theta2);
			Vector3D refractedDirection = incidentDirection.mult(refractiveIndicesRatio).sub(normal.mult(value));
			return refractedDirection;
		}
	}

	protected double[] diffuseShading(Colour lightColour, ColourCoeffs objCoeffs, 
			Vector3D normal, Vector3D lightFromPoint ){

		double diffuseRed = 0, diffuseGreen = 0, diffuseBlue = 0;
		double cosTheta = normal.dot(lightFromPoint);

		if (cosTheta > 0){	// the light is in front of the surface
			diffuseRed = lightColour.get(Params.X)*objCoeffs.getDiffuseCoeffs(Params.X)*cosTheta;
			diffuseGreen = lightColour.get(Params.Y)*objCoeffs.getDiffuseCoeffs(Params.Y)*cosTheta;
			diffuseBlue = lightColour.get(Params.Z)*objCoeffs.getDiffuseCoeffs(Params.Z)*cosTheta;
		}

		return new double[]{diffuseRed, diffuseGreen, diffuseBlue};
	}

	protected double[] specularShading(Colour lightColour, ColourCoeffs objCoeffs, 
			Vector3D normal, Vector3D lightFromPoint, Vector3D viewerFromPoint){

		double specularRed = 0, specularGreen = 0, specularBlue = 0;
		double cosAlpha = normal.mult(2*lightFromPoint.dot(normal)).sub(lightFromPoint).dot(viewerFromPoint);

		if (cosAlpha > 0){
			specularRed = lightColour.get(Params.X)*objCoeffs.getSpecularCoeffs(Params.X)*
					Math.pow(cosAlpha, objCoeffs.shininess);
			specularGreen = lightColour.get(Params.Y)*objCoeffs.getSpecularCoeffs(Params.Y)*
					Math.pow(cosAlpha, objCoeffs.shininess);
			specularBlue = lightColour.get(Params.Z)*objCoeffs.getSpecularCoeffs(Params.Z)*
					Math.pow(cosAlpha, objCoeffs.shininess);
		}

		return new double[]{specularRed, specularGreen, specularBlue};
	}

	protected boolean inShadow(IntersectionInfo intersection, LightSource light){
		Vector3D lightFromPoint = light.position.sub(intersection.position);
		Ray shadowRay = new Ray(intersection.position, lightFromPoint);

		// find the closest intersection
		IntersectionInfo shadowInfo = scene.getIntersection(shadowRay);
		if (shadowInfo != null){
			if (shadowInfo.t > Params.EPILSON &&  Math.sqrt(lightFromPoint.norm())-shadowInfo.t > Params.EPILSON){
				return true;
			}
		}/*
		for (SceneObject object : scene.sceneObjects){
			IntersectionInfo shadowInfo = object.getIntersection(shadowRay);
			if (shadowInfo != null && shadowInfo.t > Params.EPILSON){
				if (Math.pow(shadowInfo.t, 2) < lightFromPoint.norm()){
					return true;
				}
			}
		}*/
		return false;
	}	
}
